# -*- coding: utf-8 -*-
"""
数据管理模块
用于加载和管理测试数据集
"""

import json
import os
import sqlite3
from typing import Dict, List, Optional, Tuple, Any
from loguru import logger
from config import TEST_DATA_CONFIG, DATABASE_CONFIG


class DataManager:
    """数据管理器类"""
    
    def __init__(self):
        """
        初始化数据管理器
        """
        self.test_data_path = TEST_DATA_CONFIG["test_data_path"]
        self.sample_size = TEST_DATA_CONFIG["sample_size"]
        self.categories = TEST_DATA_CONFIG["categories"]
        
        # 确保数据目录存在
        os.makedirs(os.path.dirname(self.test_data_path), exist_ok=True)
    
    def load_test_data(self) -> List[Dict[str, Any]]:
        """
        加载测试数据
        
        Returns:
            List[Dict]: 测试数据列表
        """
        if os.path.exists(self.test_data_path):
            try:
                with open(self.test_data_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                logger.info(f"✅ 从文件加载了 {len(data)} 条测试数据")
                return data
            except Exception as e:
                logger.error(f"❌ 加载测试数据失败: {e}")
        
        # 如果文件不存在，从数据库生成测试数据
        logger.info("📊 从数据库生成测试数据...")
        return self._generate_test_data_from_db()
    
    def _generate_test_data_from_db(self) -> List[Dict[str, Any]]:
        """
        从数据库生成测试数据
        
        Returns:
            List[Dict]: 测试数据列表
        """
        test_data = []
        
        try:
            # 连接数据库
            conn = sqlite3.connect(DATABASE_CONFIG["path"])
            cursor = conn.cursor()
            
            # 为每个分类生成测试数据
            for category in self.categories:
                category_data = self._generate_category_test_data(cursor, category)
                test_data.extend(category_data)
            
            conn.close()
            
            # 保存测试数据到文件
            self._save_test_data(test_data)
            
            logger.info(f"✅ 生成了 {len(test_data)} 条测试数据")
            return test_data
            
        except Exception as e:
            logger.error(f"❌ 从数据库生成测试数据失败: {e}")
            return []
    
    def _generate_category_test_data(self, cursor, category: str) -> List[Dict[str, Any]]:
        """
        为特定分类生成测试数据
        
        Args:
            cursor: 数据库游标
            category: 分类名称
        
        Returns:
            List[Dict]: 该分类的测试数据
        """
        category_data = []
        
        try:
            # 查询该分类的问题和答案
            query = """
            SELECT id, question, answer, category_id 
            FROM questions 
            WHERE category_id = ? 
            ORDER BY RANDOM() 
            LIMIT ?
            """
            
            cursor.execute(query, (category, self.sample_size))
            rows = cursor.fetchall()
            
            for row in rows:
                question_id, question, answer, category_id = row
                
                # 解析答案JSON
                try:
                    if answer:
                        answer_data = json.loads(answer)
                        answer_text = answer_data.get('answer', '')
                    else:
                        answer_text = ''
                except:
                    answer_text = answer or ''
                
                # 构建测试数据项
                test_item = {
                    "id": question_id,
                    "question": question,
                    "answer": answer_text,
                    "category": category_id,
                    "keywords": self._extract_keywords_from_text(question),
                    "entities": self._extract_entities_from_text(question),
                    "difficulty": self._estimate_difficulty(question),
                    "type": self._classify_question_type(question)
                }
                
                category_data.append(test_item)
            
            logger.info(f"📋 为分类 {category} 生成了 {len(category_data)} 条测试数据")
            
        except Exception as e:
            logger.error(f"❌ 为分类 {category} 生成测试数据失败: {e}")
        
        return category_data
    
    def _extract_keywords_from_text(self, text: str) -> List[str]:
        """
        从文本中提取关键词（简单实现）
        
        Args:
            text: 输入文本
        
        Returns:
            List[str]: 关键词列表
        """
        # 简单的关键词提取（可以后续用LLM优化）
        import re
        
        # 移除标点符号，分词
        words = re.findall(r'\b\w+\b', text.lower())
        
        # 过滤停用词（简单版本）
        stop_words = {'的', '是', '在', '有', '和', '与', '或', '但', '如果', '那么', '什么', '怎么', '为什么', '如何'}
        keywords = [word for word in words if word not in stop_words and len(word) > 1]
        
        return keywords[:5]  # 返回前5个关键词
    
    def _extract_entities_from_text(self, text: str) -> List[str]:
        """
        从文本中提取实体（简单实现）
        
        Args:
            text: 输入文本
        
        Returns:
            List[str]: 实体列表
        """
        # 简单的实体提取（可以后续用NER模型优化）
        import re
        
        entities = []
        
        # 提取可能的地名、人名、机构名等
        # 这里使用简单的规则，实际应用中可以使用更复杂的NER模型
        
        # 提取大写开头的词组
        capitalized_words = re.findall(r'\b[A-Z][a-z]+\b', text)
        entities.extend(capitalized_words)
        
        # 提取数字
        numbers = re.findall(r'\d+', text)
        entities.extend(numbers)
        
        return list(set(entities))  # 去重
    
    def _estimate_difficulty(self, question: str) -> str:
        """
        估算问题难度
        
        Args:
            question: 问题文本
        
        Returns:
            str: 难度等级 (easy/medium/hard)
        """
        # 简单的难度估算规则
        question_len = len(question)
        
        # 检查复杂词汇
        complex_words = ['分析', '评估', '比较', '解释', '推理', '计算', '设计', '优化']
        has_complex_words = any(word in question for word in complex_words)
        
        if question_len < 20 and not has_complex_words:
            return 'easy'
        elif question_len > 50 or has_complex_words:
            return 'hard'
        else:
            return 'medium'
    
    def _classify_question_type(self, question: str) -> str:
        """
        分类问题类型
        
        Args:
            question: 问题文本
        
        Returns:
            str: 问题类型
        """
        question_lower = question.lower()
        
        if any(word in question_lower for word in ['什么', 'what', '是什么']):
            return 'factual'  # 事实性问题
        elif any(word in question_lower for word in ['怎么', 'how', '如何']):
            return 'procedural'  # 程序性问题
        elif any(word in question_lower for word in ['为什么', 'why', '原因']):
            return 'causal'  # 因果性问题
        elif any(word in question_lower for word in ['比较', 'compare', '区别']):
            return 'comparative'  # 比较性问题
        else:
            return 'general'  # 一般性问题
    
    def _save_test_data(self, test_data: List[Dict[str, Any]]) -> None:
        """
        保存测试数据到文件
        
        Args:
            test_data: 测试数据列表
        """
        try:
            with open(self.test_data_path, 'w', encoding='utf-8') as f:
                json.dump(test_data, f, ensure_ascii=False, indent=2)
            logger.info(f"✅ 测试数据已保存到 {self.test_data_path}")
        except Exception as e:
            logger.error(f"❌ 保存测试数据失败: {e}")
    
    def get_questions_by_category(self, category: str) -> List[Dict[str, Any]]:
        """
        获取指定分类的问题
        
        Args:
            category: 分类名称
        
        Returns:
            List[Dict]: 该分类的问题列表
        """
        test_data = self.load_test_data()
        return [item for item in test_data if item.get('category') == category]
    
    def get_questions_by_type(self, question_type: str) -> List[Dict[str, Any]]:
        """
        获取指定类型的问题
        
        Args:
            question_type: 问题类型
        
        Returns:
            List[Dict]: 该类型的问题列表
        """
        test_data = self.load_test_data()
        return [item for item in test_data if item.get('type') == question_type]
    
    def get_questions_by_difficulty(self, difficulty: str) -> List[Dict[str, Any]]:
        """
        获取指定难度的问题
        
        Args:
            difficulty: 难度等级
        
        Returns:
            List[Dict]: 该难度的问题列表
        """
        test_data = self.load_test_data()
        return [item for item in test_data if item.get('difficulty') == difficulty]
    
    def create_knowledge_base(self) -> Dict[str, List[str]]:
        """
        创建知识库（用于知识图谱检索）
        
        Returns:
            Dict: 知识库字典
        """
        test_data = self.load_test_data()
        knowledge_base = {
            'entities': [],
            'relations': [],
            'facts': []
        }
        
        for item in test_data:
            # 收集实体
            knowledge_base['entities'].extend(item.get('entities', []))
            
            # 构建事实
            fact = f"{item['question']} -> {item['answer']}"
            knowledge_base['facts'].append(fact)
        
        # 去重
        knowledge_base['entities'] = list(set(knowledge_base['entities']))
        
        return knowledge_base


# 全局数据管理器实例
data_manager = DataManager()


if __name__ == "__main__":
    # 测试数据管理器
    dm = DataManager()
    
    # 加载测试数据
    test_data = dm.load_test_data()
    print(f"加载了 {len(test_data)} 条测试数据")
    
    if test_data:
        print(f"示例数据: {test_data[0]}")
        
        # 测试分类查询
        category_data = dm.get_questions_by_category('1')
        print(f"分类1的问题数量: {len(category_data)}")
        
        # 创建知识库
        kb = dm.create_knowledge_base()
        print(f"知识库实体数量: {len(kb['entities'])}")
        print(f"知识库事实数量: {len(kb['facts'])}")